/*
 * Copyright (c) Martin Kinkelin
 *
 * See the "License.txt" file in the root directory for infos
 * about permitted and prohibited uses of this code.
 */

/*
 * Very simple Windows service wrapping a non-service executable.
 *
 * Executing the program using the
 * a) "-i" switch installs the service,
 * b) "-u" switch uninstalls the service.
 *
 * The service will be started automatically when the system is started
 * using the LocalSystem account.
 * All it does is starting a child process and terminating when either
 * a) the child process has exited or
 * b) the service is stopped (terminating the child process).
 *
 * Modify the 4 macros below to use the service for another executable.
 */

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <tchar.h>

#include <iostream>
#include <string>

using namespace std;



#define NAME "RsyncServer" // must not contain whitespaces
#define DISPLAY_NAME "Rsync server"
#define DESCRIPTION "Rsync server for backup/synchronization."

// Command used to start the child process.
// "%0" is replaced by this program's directory,
// "%1" by the corresponding Cygwin path.
#define COMMAND "\"%0\\Rsync\\rsync.exe\" --daemon --no-detach --config=\"%1/rsyncd.conf\""



/// <summary>Entry point for the service.</summary>
void WINAPI ServiceMain(__in DWORD dwArgc, __in LPTSTR* lpszArgv);

/// <summary>
/// Control handler function of the service.
/// This function reacts to start/stop requests etc.
/// </summary>
DWORD WINAPI ServiceCtrlHandlerEx(__in DWORD dwControl, __in DWORD dwEventType,
                                  __in LPVOID lpEventData, __in LPVOID lpContext);

/// <summary>
/// Returns the command, using the COMMAND macro as special format string.
/// Make sure you free() this string.
/// </summary>
TCHAR *GetCommand();



/// <summary>Entry point for the program.</summary>
int main(int argc, char* argv[])
{
	// try to install the service if the "-i" switch is used
	if (argc > 1 && _tcsicmp(argv[1], "-i") == 0)
	{
		TCHAR buffer[MAX_PATH + 2];

		// try connect to the local service control manager
		SC_HANDLE scm = OpenSCManager(NULL, NULL, SC_MANAGER_CREATE_SERVICE);
		if (scm == NULL)
		{
			cerr << "Could not connect to the service control manager." << endl;
			return 1;
		}

		// get the path of this program in double quotes
		buffer[0] = '\"';
		GetModuleFileName(NULL, buffer + 1, MAX_PATH);
		_tcscat_s(buffer, sizeof(buffer), "\"");

		// try to install the service
		SC_HANDLE service = CreateService(scm, NAME, DISPLAY_NAME, SERVICE_CHANGE_CONFIG,
			SERVICE_WIN32_OWN_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL,
			buffer, NULL, NULL, NULL, NULL, "");
		if (service == NULL)
		{
			cerr << "Could not install the service." << endl;
			CloseServiceHandle(scm);
			return 2;
		}

		// attempt to set the description
		SERVICE_DESCRIPTION description;
		description.lpDescription = DESCRIPTION;
		ChangeServiceConfig2(service, SERVICE_CONFIG_DESCRIPTION, &description);

		// attempt to enable the autostart delay supported by Vista
		SERVICE_DELAYED_AUTO_START_INFO delayedAutoStartInfo;
		delayedAutoStartInfo.fDelayedAutostart = TRUE;
		ChangeServiceConfig2(service, SERVICE_CONFIG_DELAYED_AUTO_START_INFO, &delayedAutoStartInfo);

		cout << "The service has been installed." << endl;

		CloseServiceHandle(service);
		CloseServiceHandle(scm);

		return 0;
	}

	// uninstall the service if the "-u" switch is used
	if (argc > 1 && _tcsicmp(argv[1], "-u") == 0)
	{
		// try to connect to the local service control manager
		SC_HANDLE scm = OpenSCManager(NULL, NULL, SC_MANAGER_CONNECT);
		if (scm == NULL)
		{
			cerr << "Could not connect to the service control manager." << endl;
			return 1;
		}

		SC_HANDLE service = OpenService(scm, NAME, SERVICE_STOP | DELETE);
		if (service == NULL)
		{
			cerr << "Could not uninstall the service." << endl;
			CloseServiceHandle(scm);
			return 3;
		}

		// try to stop and delete the service
		SERVICE_STATUS serviceStatus;
		if ((!ControlService(service, SERVICE_CONTROL_STOP, &serviceStatus) && GetLastError() != ERROR_SERVICE_NOT_ACTIVE) ||
		    !DeleteService(service))
		{
			cerr << "Could not uninstall the service." << endl;
			CloseServiceHandle(service);
			CloseServiceHandle(scm);
			return 3;
		}

		cout << "The service has been uninstalled." << endl;

		CloseServiceHandle(service);
		CloseServiceHandle(scm);

		return 0;
	}

	// try to start the actual service
	SERVICE_TABLE_ENTRY dispatchTable[] = { { NAME, ServiceMain }, {0, 0} };
	if (!StartServiceCtrlDispatcher(dispatchTable))
		return 10;

	return 0;
}



SERVICE_STATUS _status;
SERVICE_STATUS_HANDLE _statusHandle;

bool _isRunning = true;


void WINAPI ServiceMain(__in DWORD dwArgc, __in LPTSTR* lpszArgv)
{
	// try to register the service control handler
	_statusHandle = RegisterServiceCtrlHandlerEx(NAME, ServiceCtrlHandlerEx, NULL);
	if (_statusHandle == NULL)
		return;

	TCHAR *command = GetCommand();

	DWORD creationFlags = CREATE_NO_WINDOW;
#ifdef UNICODE
	creationFlags |= CREATE_UNICODE_ENVIRONMENT;
#endif

	STARTUPINFO startupInfo;
	memset(&startupInfo, 0, sizeof(startupInfo));
	startupInfo.cb = sizeof(startupInfo);

	// set the "CYGWIN" environment variable to "nontsec" to
	// disable modifying the permissions of files and folders
	// the child process will inherit this process' environment
	SetEnvironmentVariable("CYGWIN", "nontsec");

	// try to create the child process
	PROCESS_INFORMATION process;
	if (!CreateProcess(NULL, command, NULL, NULL, FALSE, creationFlags,
		NULL, NULL, &startupInfo, &process))
	{
		free(command);
		return;
	}

	free(command);

	_status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
	_status.dwCurrentState = SERVICE_RUNNING;
	_status.dwControlsAccepted = SERVICE_ACCEPT_STOP;
	_status.dwWin32ExitCode = NO_ERROR;
	_status.dwServiceSpecificExitCode = 0;
	_status.dwCheckPoint = 0;
	_status.dwWaitHint = 0;

	if (!SetServiceStatus(_statusHandle, &_status))
		return;

	DWORD exitCode = STILL_ACTIVE;

	// loop as long as the service has not been stopped and
	// the child process is still running
	while (_isRunning && exitCode == STILL_ACTIVE)
	{
		Sleep(1000);

		// update the exit code
		GetExitCodeProcess(process.hProcess, &exitCode);
  }

	// terminate the child process if still running
	if (exitCode == STILL_ACTIVE)
		TerminateProcess(process.hProcess, NO_ERROR);
	else
		_status.dwWin32ExitCode = exitCode;

	CloseHandle(process.hThread);
	CloseHandle(process.hProcess);

	_status.dwCurrentState = SERVICE_STOPPED; 

	SetServiceStatus(_statusHandle, &_status);
} 

DWORD WINAPI ServiceCtrlHandlerEx(__in DWORD dwControl, __in DWORD dwEventType,
                                  __in LPVOID lpEventData, __in LPVOID lpContext)
{ 
	switch (dwControl)
	{
		case SERVICE_CONTROL_STOP:
			_isRunning = false;
			_status.dwCurrentState = SERVICE_STOP_PENDING;
			SetServiceStatus(_statusHandle, &_status);
			return NO_ERROR;

		// the MSDN suggests returning NO_ERROR even if this event is not handled:
		case SERVICE_CONTROL_INTERROGATE:
			return NO_ERROR;
	}

	return ERROR_CALL_NOT_IMPLEMENTED;
} 

TCHAR *GetCommand()
{
	// compute the directory of this program
	TCHAR rawDir[MAX_PATH];
	GetModuleFileName(NULL, rawDir, MAX_PATH);
	string dir(rawDir);
	dir.erase(dir.find_last_of('\\'));

	// compute the corresponding Cygwin path
	string::size_type i = 0;
	string cygwinDir;
	if (dir[1] == ':')
	{
		// local path, not a NetBIOS path
		cygwinDir = "/cygdrive/";
		cygwinDir += dir[0];
		i = 2;
	}
	for (; i < dir.length(); ++i)
	{
		cygwinDir += (dir[i] == '\\' ? '/' : dir[i]);
	}

	// split the COMMAND format string in tokens separated by '%'
	TCHAR *commandFormat = _tcsdup(COMMAND); // the format string will be modified, a literal string cannot be used
	TCHAR *token, *nextToken = NULL;
	token = _tcstok_s(commandFormat, "%", &nextToken); // get the first token

	string commandLine(token);

	// process the remaining tokens
	while ((token = _tcstok_s(NULL, "%", &nextToken)) != NULL)
	{
		// "%0" in the format string is replaced by dir, "%1" by cygwinDir
		if (token[0] == '1')
			commandLine += cygwinDir;
		else
			commandLine += dir;

		// append the rest of the token
		commandLine += token + 1;
	}

	free(commandFormat);

	return _tcsdup(commandLine.c_str());
}
